#    Authors:    Chao Li, Handing Wang, Jun Zhang, Wen Yao, Tingsong Jiang
#    Xidian University, China
#    Defense Innovation Institute, Chinese Academy of Military Science, China
#    EMAIL:      lichaoedu@126.com, hdwang@xidian.edu.cn
#    DATE:       February 2022
# ------------------------------------------------------------------------
# This code is part of the program that produces the results in the following paper:
#
# Chao Li, Handing Wang, Jun Zhang, Wen Yao, Tingsong Jiang, An Approximated Gradient Sign Method Using Differential Evolution For Black-box Adversarial Attack, IEEE Transactions on Evolutionary Computation, 2022.
#
# You are free to use it for non-commercial purposes. However, we do not offer any forms of guanrantee or warranty associated with the code. We would appreciate your acknowledgement.
# ------------------------------------------------------------------------


import numpy as np
import torch
import torch.utils.data as Data
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from FNSack.FNS_DE import FNSDE

class FNSack:

    def __init__(self, config, neter, dataer):

        self.config = config
        self.neter = neter
        self.model = self.neter.net
        self.dataer = dataer
        ## self.test_loader = test_loader
        self.device = self.config['device']
        self.isTarget = False
        self.class_index = None
    
    def set_attack_mode(self, isTarget=False, class_index=None):
        
        self.isTarget = isTarget

        if self.isTarget:
            print('Change into target attack mode')
            if class_index == None:
                raise Exception('FNS target attack mode, but the class is None!')
        
            self.class_index = class_index

    def select_second_labels_images(self, second_label, first_label):
        sum_images = []
        
        if self.dataer.dataset_name == 'ImageNet':
            test_loader = self.dataer.get_class_loader(class_index=second_label, isTrain=False, batch_size=1, isThreePart=True)
        else:
            test_loader = self.dataer.get_class_loader(class_index=second_label, isTrain=False, batch_size=1)

        for images, labels in test_loader:
            if labels == second_label:
                images = images.to(self.device)
                outputs = self.model(images)
                outputs = outputs.cpu().detach().numpy()
                if np.argmax(outputs) == second_label:
                    min_value=np.min(outputs)
                    outputs.itemset(second_label, min_value)
                    if np.argmax(outputs) == first_label:
                        target_images = torch.tensor(images)
                        target_images = np.array(target_images.cpu())
                        sum_images.append(target_images)
        sum_images = np.array(sum_images)
        num_images = np.size(sum_images, 0)
        if num_images > 1:
            sum_images = sum_images.squeeze()
        else:
            sum_images = sum_images.squeeze()
            sum_images = torch.tensor(sum_images)
            sum_images = sum_images.unsqueeze(0)
            sum_images = sum_images.detach().numpy()

        if self.config['dataset'] == 'MNIST' or self.config['dataset'] == 'FMNIST':

            sum_images = torch.tensor(sum_images)
            sum_images = sum_images.unsqueeze(1)
            sum_images = sum_images.detach().numpy()

        return sum_images

    def FNS_attack(
        self,
        images,
        labels,
    ):

        ## 被注释掉的是源代码实现test_loader的方式，现在通过接口改写.
        # batch_size = 1
        # test_dataset = dsets.CIFAR10(root='/home/amax/文档/LICHAO_code/lc_code_CIFAR10_Non_target/CIFAR_data',
        #                                 download=False,
        #                                 transform=transforms.Compose([
        #                                     transforms.ToTensor(),
        #                                     transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        #                                             ]),
        #                                 train=False)
        # test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
        #                                         batch_size=batch_size,
        #                                         shuffle=True)


        self.model.eval()

        FDEer = FNSDE(self.neter, self.config, isTarget=self.isTarget, class_index=self.class_index)

        images = images.to(self.device)
        labels = labels.to(self.device)
        outputs = self.model(images)
        _, pre = torch.max(outputs.data, 1)

        outputs = outputs.cpu().detach().numpy()
        min_value = np.min(outputs)
        outputs.itemset(labels, min_value)
        second_label = np.argmax(outputs)
        sum_images = self.select_second_labels_images(second_label, labels)

        images = FDEer.FDE(images,sum_images, second_label, labels)
        
        images = images.to(self.device)
        labels = labels.to(self.device)
        outputs = self.model(images)
        _, pre = torch.max(outputs.data, 1)

        ## untarget atttack
        if self.isTarget == False:

            if pre != labels:
                return True, images
            else:
                return False, images
        ## target attack
        else:
            
            if pre[0] == self.class_index:
                return True, images
            else:
                return False, images